import json
import os.path
import platform
import shutil

import requests
from PyQt6.QtGui import QDragEnterEvent, QDropEvent
from PyQt6.QtWidgets import QPushButton, QWidget, QTextEdit, \
    QSizePolicy, QLabel, QScrollArea, QCheckBox, QLineEdit
from PyQt6.uic import loadUi

import global_config
import main_model
import midjourney_config
from global_config import *
from scripts import db_utils, beans
from scripts.configui.ImageGenerateTask import ImageTask, SVDTask, GenerateManager

from scripts.jianying import draft_parser
from ui.adapters.ScrollviewAdapter import ScrollviewAdapter
from ui.dialogs.CommonPromptsDialog import CommonPromptsDialog
from ui.model import DrawPageModel
from ui.widgets.notice_dialog import ConfirmDialog


# 原创绘制页面
class DrawPage(QWidget):
    def __init__(self):
        super().__init__()
        self.selectAll = False
        self.prompt_config_btn: QPushButton = None
        self.menu_btn: QPushButton = None
        self.save_btn: QPushButton = None
        self.loading: QLabel = None
        self.global_shuff_cb: QCheckBox = None
        self.parse_btn: QPushButton = None
        self.svd_btn: QPushButton = None
        self.jianying_btn: QPushButton = None
        self.generate_btn: QPushButton = None
        self.reasoning_btn: QPushButton = None
        self.global_shuff_tv: QLineEdit = None
        self.storyboards: beans.Storyboard = None
        self.scrollview: QScrollArea = None
        self.listAdapter: ScrollviewAdapter = None
        self.videoDialog = None
        self.initView()
        self.imgGenManager = GenerateManager(self.on_generate_images)
        self.imgGenManager.startTask()

    def initView(self):
        loadUi('./source/draw_page.ui', self)
        self.listAdapter = ScrollviewAdapter(self.scrollview, callback=self.on_adapter_item_callback)
        self.loading.hide()
        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)

        self.video = get_video_path()

        self.menu_btn.clicked.connect(self.on_open_dir)
        self.save_btn.clicked.connect(self.update_workflow)
        self.generate_btn.clicked.connect(self.generate_checked_item)
        self.reasoning_btn.clicked.connect(lambda: self.on_adapter_item_callback('推理', -1))
        self.prompt_config_btn.clicked.connect(self.open_prompt_dialog)
        self.select_cb.clicked.connect(self.select_all_holder)
        self.svd_btn.clicked.connect(self.generate_svd)
        self.jianying_btn.clicked.connect(self.generate_jianying_draft)

        # 加载数据
        storyboards_path = get_storyboards_path()
        config_json_path = get_config_json()
        if os.path.exists(storyboards_path):
            self.storyboards = db_utils.read_workflow(storyboards_path)
            self.load_storyboards(self.storyboards)
        self.global_shuff_tv.setText(db_utils.read_config_json(config_json_path))

    def select_all_holder(self):
        self.selectAll = False if self.selectAll else True
        for holder in self.listAdapter.itemholders:
            holder.setSelected(self.selectAll)

    # 加载解析完的 分镜数据
    def load_storyboards(self, storyboards):
        self.listAdapter.load(storyboards)
        # 创建一个信号，触发时传递当前时间给槽函数

    def on_open_dir(self):
        if platform.system() == 'Darwin':
            os.system(f'open /{workflows_dir}{get_cur_project()}')
        else:
            os.startfile(f'{workflows_dir}{get_cur_project()}')

    # 打开提示词配置弹出
    def open_prompt_dialog(self):
        CommonPromptsDialog(None).get_result()

    # 推理提示词
    def on_adapter_item_callback(self, command, index):
        if command == '推理':
            if index >= 0:
                def on_call(result):
                    comid, content = result
                    if content is None:return
                    if len(content) > 0:
                        self.listAdapter.itemholders[index].bean.promots = content[0]
                        self.listAdapter.itemholders[index].update_imagesUI()
                DrawPageModel.reasoning(index, self.listAdapter.itemholders,on_call)

            else:
                selected_holders = []
                for v in self.listAdapter.itemholders:
                    if v.index_checkbox.isChecked():
                        selected_holders.append(v)

                def on_call_muti(content):
                    comid, content = content
                    if content is None:return
                    for i in range(0, len(selected_holders)):
                        if i < len(content):
                            selected_holders[i].bean.promots = content[i]
                            selected_holders[i].update_imagesUI()
                DrawPageModel.reasoning(index, selected_holders,on_call_muti)

            db_utils.save_wordflow(get_storyboards_path(), self.storyboards)

    # 生成选中的分镜图片
    def generate_checked_item(self):
        for holder in self.listAdapter.itemholders:
            if holder.index_checkbox.isChecked():
                holder.setSelected(False)
                if holder.bean.ref_img is None and holder.bean.local_img is not None:
                    print("upload")
                    res = requests.post(url=comfyui_upload,
                                        files={'image': open(holder.bean.local_img, 'rb')})
                    holder.bean.ref_img = res.json()['name']
                    db_utils.save_wordflow(get_storyboards_path(), self.storyboards)

                qp_config = db_utils.load_project_config()
                promots = holder.bean.to_full_prompts(qp_config['positive_prompt'], self.global_shuff_tv.text())
                # self.loading.show()
                file = holder.bean.ref_img
                # if not os.path.exists(file):
                if file is not None:
                    self.imgGenManager.enQueue(
                        ImageTask(promots, qp_config['negative_prompt'], holder.index, photoname=file,
                                  width=qp_config['width'], height=qp_config['height']))
                else:
                    self.imgGenManager.enQueue(ImageTask(promots, qp_config['negative_prompt'], holder.index,
                                                         width=qp_config['width'], height=qp_config['height']))

    # 图片转视频
    def generate_svd(self):
        ConfirmDialog('提示', f'图片转视频，禁用中。联系作者').get_result()
        #
        # for holder in self.listAdapter.itemholders:
        #     if holder.index_checkbox.isChecked():
        #         imgsdir = f'{get_workflows_img_dir()}\\{holder.index}\\'
        #         if not os.path.exists(imgsdir):
        #             continue
        #         files = os.listdir(imgsdir)
        #         size = len(files)
        #         for i in range(0, size):
        #             self.imgGenManager.enQueue(SVDTask(holder.index, comfyui_output + files[i], i))

    # 生成剪映草稿
    def generate_jianying_draft(self):
        draft_parser.sycn_draft()

    def update_workflow(self):
        storyboards_path = get_storyboards_path()
        db_utils.save_config_json(get_config_json(), self.global_shuff_tv.text())
        if self.storyboards is not None:
            db_utils.save_wordflow(storyboards_path, self.storyboards)

    def on_generate_images(self, success, result_images, taskdir, index, isEnd, isVideo):
        if isEnd:
            self.loading.hide()
        if not success: return

        print(index, result_images)
        for holder in self.listAdapter.itemholders:
            if holder.index == index:
                # 自动选中第一个图片
                if holder.bean.select_img is None or not os.path.exists(holder.bean.select_img):
                    aiImgDir = f'{global_config.get_workflows_img_dir()}\\{index}\\'
                    if os.path.exists(aiImgDir):
                        filelist = os.listdir(aiImgDir)
                        holder.bean.select_img = f'{aiImgDir}{filelist[0]}' if len(filelist) > 0 else None
                db_utils.save_wordflow(get_storyboards_path(), self.storyboards)
                holder.update_imagesUI()

        if isVideo:
            if self.videoDialog == None:
                self.videoDialog = ConfirmDialog('提示', f'分镜{index}，视频生成完毕，打开目录查看')
            self.videoDialog.message_label.setText(f'分镜{index}，视频生成完毕，打开目录查看')
            self.videoDialog.get_result()

    def resizeEvent(self, event):
        # 获取新的窗口尺寸
        size = self.size()
        print(f"窗口尺寸变化为宽度: {size.width()}， 高度: {size.height()}", event)
        geo = self.loading.geometry()
        geo.setWidth(size.width())
        geo.setHeight(size.height())
        self.loading.setGeometry(geo)
